#!usr/bin/env python3
# encoding:utf-8
# 使用imageai进行图片的分类的预测，模型一共训练了三个分类 1、normal 2、anchor 3、boundry 
# 如果非anchor和boundry的图片，则认为其没有信息；可以按目录来查询，也可以按图片进行查询；
from __future__ import division
 
"""
__Author__:云起
功能： Python基于ImageAI实现完成的流程：数据集构建、模型训练、识别预测
"""
import os,sys
import threading
from pathlib import Path
import shutil
import torch
import torchvision.transforms as transforms
import torchvision
from PIL import Image
os.path.join(os.path.dirname(__file__), '../')
sys.path.append(os.path.join(os.path.dirname(__file__), '../'))
from public import newsCNN

mean_nums = [0.485, 0.456, 0.406]
std_nums = [0.229, 0.224, 0.225]

normalize = transforms.Normalize(mean_nums, std_nums)
transform = transforms.Compose([
    transforms.Resize((64,64)),

    transforms.ToTensor(), normalize]
)

classes = ["anchor","boundary","normal"]


class imgCategory:

    #初始化数据模型
    def __init__(self):

        model_path = "./model/news_classification_20220426.pkl"
        self.news_model = torch.load(model_path)
        # 如果有GPU卡，就需要加上下面执行语句，把模型加载到GPU中;
        # self.news_model.eval()

    def imgPredict(self, pic_path):

        img = Image.open(pic_path)
        print(img.size)
        # 对图像进行归一化
        img_p = transform(img)
        # print(img_p.shape)
        
        # 增加一个维度, 因为图片转为3个通道后，只有三个维度，而模型需要四个维度，所以增加一个维度进行处理;
        img_normalize = torch.unsqueeze(img_p, 0)
        # print(img_normalize.shape)
        #调用模型进行预测
        pred=self.news_model(img_normalize)
        #预测结果进行排序
        _, indices = torch.sort(pred, descending=True)
        #处理相应的匹配分值
        percentage = torch.nn.functional.softmax(pred, dim=1)[0] * 100
        #获取相应的分类及分值数组
        prediction = [[classes[idx], percentage[idx].item()] for idx in indices[0][:5]]
        # print(prediction)
        
        ret = "normal"
        for i in prediction:
            #如果数值有超50的，则认为其归属于该分类
            if(i[1] > 50):
                ret = i[0]
                break
        return ret

    def pathPredict(self, img_path):

        #准备一个临时目录存放关键帧转场的文件
        tmp_pstr = img_path + "/category"
        tmp_path = Path(tmp_pstr)
        tmp_path.mkdir(parents=True, exist_ok=True)
        #get  path files list;
        dlist = os.listdir(img_path)
        #只获取列表中的.jpg内容；
        nlist = []
        for m in  range(len(dlist)):
            if dlist[m].endswith('.jpg'):
                nlist.append(dlist[m])

        img_list =sorted(nlist)    #文件名按字母排序
        img_nums =len(img_list)
        for i in range(img_nums):
            pimg = img_path + "/" +img_list[i]
            class_name = self.imgPredict(pimg)
            #先去掉原文件中的后缀，重新组合文件名
            n_f = img_list[i][:-4]
            ppimg = tmp_pstr+"/"+n_f+"_"+class_name+".jpg"
            shutil.copyfile(pimg, ppimg)
        
        return self.getPathImgList(tmp_pstr)

    #获取路径图片的列表
    def getPathImgList(self, img_path):
        
        dlist = os.listdir(img_path)
        #只获取列表中的.jpg内容；
        nlist = []
        for m in  range(len(dlist)):
            if dlist[m].endswith('.jpg'):
                nlist.append(dlist[m])

        return nlist
    
if __name__=='__main__':

    ic = imgCategory()
    # for i in range(1,58):
    #     print(i)
    #     ic.pathPredict("/tmp/autostrip/"+str(i))
    ic.pathPredict("/root/newspower/dataset/new_data/train/anchor")
    ic.pathPredict("/root/newspower/dataset/new_data/train/boundary")
    ic.pathPredict("/root/newspower/dataset/new_data/train/normal")
